package com.wang.tool.trace;

import com.wang.tool.util.SimpleLogUtil;

import java.util.Map;

public class TraceHolder {

    public static final String TRACE_ID_HEADER = "X-Trace-Id";
    public static final String RPC_ID_HEADER = "X-Rpc-Id";
    public static final String ATTRIBUTES_HEADER = "X-Attributes-Id";


    private static TraceTransmitter transmitter;

    private static TraceBuilder traceBuilder;


    public static TraceContext getCurrent() {
        if (transmitter == null) {
            return null;
        }
        try {
            return transmitter.get();
        } catch (Throwable e) {
            SimpleLogUtil.error("fetchTraceError", e);
            return null;
        }
    }


    public static String getTraceId() {
        TraceContext current = getCurrent();
        return current == null ? null : current.getTraceId();
    }

    public static String getRpcId() {
        TraceContext current = getCurrent();
        return current == null ? null : current.getRpcId();
    }

    public static String getReqId() {
        TraceContext current = getCurrent();
        return current == null ? null : current.buildReqId();
    }

    public static Map<String,Object> getAttributes() {
        TraceContext current = getCurrent();
        return current == null ? null : current.getAttributes();
    }

    public static Object getAttribute(String key) {
        if (key == null) {
            return null;
        }
        TraceContext context = getCurrent();
        return context == null ? null : context.getAttribute(key);
    }

    public static void putAttribute(String key, Object val) {
        if (key == null) {
            return;
        }
        TraceContext context = getCurrent();
        if (context == null) {
            return;
        }
        context.putAttribute(key, val);
    }


    public static void set(TraceContext context) {
        if (transmitter == null) {
            return;
        }
        try {
            transmitter.set(context);
        } catch (Throwable e) {
            SimpleLogUtil.error("setTraceError", e);
        }
    }

    public static void clear() {
        set(null);
    }


    public static TraceContext getOrCreateTrace(TraceContext current) {
        if (current != null) {
            return current;
        }
        TraceContext context = TraceContext.newTrace();
        set(context);
        return context;
    }


    public static TraceContext nexTrace(TraceContext current) {
        TraceContext context = TraceContext.nextTrace(current);
        set(context);
        return context;
    }

    public static String generateTraceId() {
        try {
            return traceBuilder == null ? null : traceBuilder.generateTrace();
        } catch (Throwable e) {
            SimpleLogUtil.error("generateTraceIdError", e);
            return null;
        }
    }


    public static void setTransmitter(TraceTransmitter transmitter) {
        TraceHolder.transmitter = transmitter;
        SimpleLogUtil.info("change traceTransmitter, " + transmitter);
    }

    public static void setTraceBuilder(TraceBuilder traceBuilder) {
        TraceHolder.traceBuilder = traceBuilder;
        SimpleLogUtil.info("change traceBuilder, " + traceBuilder);

    }
}
